#本脚本实现了将视频转化为图片，并把视频与图片的关系映射成json文件
#json文件还需优化，以便后续形成不会重复生成对应帧数的图片(待完成)
#达到新添加视频到原目录中，不会重复生成。
#存储图片名为映射序号_帧数.jpg

from pathlib import Path
import argparse
import cv2
import json
from PIL import Image
import time

def parse_arguments():
    parser = argparse.ArgumentParser(description="video2Img")
    parser.add_argument('--video_dir', type=str, default=r"D:\tempDownload\dataset\elevator\2024-01-12")  #视频路径
    parser.add_argument('--img_outdir', type=str, default=r"D:\tempDownload\dataset\elevator\mydata_img")  #图片输出路径
    parser.add_argument('--skip_fps', type=int, default=15)  #跳帧次数
    parser.add_argument('--regenerate-images', type=bool, default=False)  #是否重新生成所有视频
    return parser.parse_args()

#判断源路径和目标路径是否存在
def judgeOriginTarget(origin_file_path, target_file_path):
    origin_file_path = Path(origin_file_path)
    target_file_path = Path(target_file_path)
    if not origin_file_path.exists():
        print("you file is not exist")
        exit()

    if not target_file_path.exists():
        print(f"目标目录不存在，在{target_file_path}创建目录")
        target_file_path.mkdir(parents=True, exist_ok=True)

def main():
    #读取参数
    args = parse_arguments()
    #检查文件夹是否存在
    judgeOriginTarget(args.video_dir, args.img_outdir)

    #定义变量
    video_paths = Path(args.video_dir).glob("*.mp4")
    output_path = args.img_outdir
    skip_fps = args.skip_fps
    regenerate_images = args.regenerate_images

    #清空目标目录下的所有图片
    if regenerate_images:
        for image in Path(output_path).glob("*.jpg"):
            image.unlink()
    #将视频与对应的图片关系存储为json文件
    video_mapping_index = {}
    video_mapping_index_json_path = f"{output_path}/video_mapping_index.json"
    #若video_mapping_index_json_path 已经存在则读取
    if Path(video_mapping_index_json_path).exists():
        with open(video_mapping_index_json_path, "r", encoding="utf-8") as json_file:
            video_mapping_index = json.load(json_file)

    for video_path in video_paths:
        # #若映射关系已存在，则跳过
        if video_path.stem in video_mapping_index and not args.regenerate_images:
            print(f"{video_path.stem}已存在")
            continue
        #若文件夹中有加入新视频，则赋予新的映射关系
        video_mapping_index[video_path.stem] = len(video_mapping_index) if video_mapping_index.get(video_path.stem, -1) == -1 else video_mapping_index.get(video_path.stem)
        #开始读取视频
        vid = cv2.VideoCapture(str(video_path))
        total_frames = int(vid.get(cv2.CAP_PROP_FRAME_COUNT))  #获取视频帧数
        print(f"{video_path.stem}有{total_frames}帧")
        print(f"开始存储序号{video_mapping_index[video_path.stem]}{video_path.stem}的视频")
        for j in range(total_frames):
            #不要高速
            ret, frame = vid.read()
            if not ret:
                continue
            #结束不要高速
            if j % skip_fps == 0:
                # 开始高速跳帧，想高速可以打开这个
                # vid.set(cv2.CAP_PROP_POS_FRAMES, j)
                # ret, frame = vid.read()
                # if not ret:
                #     continue
                # 结束高速
                output_filename = f"{output_path}\\{video_mapping_index[video_path.stem]:03d}_{j:04d}.jpg"
                # rotated_frame = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)
                # cv2.imwrite(output_filename, rotated_frame) #cv2保存不了中文名
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # PIL需要RGB格式
                frame = Image.fromarray(frame)
                frame.save(output_filename)
                print(output_filename)
            time.sleep(0.01) #部分机器无法高速进行处理 sleep一会
    #保存json文件 写入json
    with open(video_mapping_index_json_path, "w", encoding="utf-8") as json_file:
        json.dump(video_mapping_index, json_file, indent=4)

if __name__ == "__main__":
    main()
